﻿// Gateway.Api 项目的版权、商标、专利和其他相关权利均受相应法律法规的保护。使用本项目应遵守相关法律法规和许可证的要求。
//
// 本项目主要遵循 MIT 许可证和 Apache 许可证（版本 2.0）进行分发和使用。许可证位于源代码树根目录中的 LICENSE-MIT 和 LICENSE-APACHE 文件。
//
// 不得利用本项目从事危害国家安全、扰乱社会秩序、侵犯他人合法权益等法律法规禁止的活动！任何基于本项目二次开发而产生的一切法律纠纷和责任，我们不承担任何责任！

using System.Diagnostics;
using Aurora.AI.Common.Constants;
using Aurora.AI.Common.Options;
using Aurora.AI.SqlSugar.Framework.Repositories;
using Aurora.AI.Wiki.Domain.Domain.Wikis.Aggregates;
using Mapster;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.KernelMemory;
using Microsoft.KernelMemory.MemoryStorage;
using SqlSugar;

namespace Aurora.AI.Wiki.Services.Wiki;

public class WikiDetailService(
    IRepository<WikiDetail> repository,
    WikiMemoryService wikiMemoryService,
    WikiSettingService wikiSettingService,
    IServiceProvider serviceProvider)
{
    public async Task<WikiDetail> Add(WikiDetail input)
    {
        WikiDetail request = new()
        {
            WikiId = input.WikiId,
            FileName = input.FileName,
            Path = input.Path,
            FileId = input.FileId,
            DataCount = 0,
            State = 0,
            Type = input.Type,
            TrainingPattern = input.TrainingPattern,
            Mode = input.Mode,
            MaxTokensPerLine = input.MaxTokensPerLine,
            MaxTokensPerParagraph = input.MaxTokensPerParagraph,
            OverlappingTokens = input.OverlappingTokens,
            QAPromptTemplate = input.QAPromptTemplate,
            CreationTime = DateTime.Now,
            Creator = input.Creator
        };
        await repository.InsertAsync(request);
        return request;
    }


    /// <summary>
    /// 更新状态 以及索引数量
    /// </summary>
    /// <param name="id"></param>
    /// <param name="state">0 默认 1 完成 2失败</param>
    /// <returns></returns>
    public async Task UpdateDetailsState(string id, int state)
    {
        int wikiDetailVectorCountAsync = await this.GetWikiDetailVectorCountAsync(id);

        await repository.GetDb().Updateable<WikiDetail>()
            .SetColumns(x => x.State == state)
            .SetColumns(x => x.ModificationTime == DateTime.Now)
            .SetColumns(x => x.DataCount == wikiDetailVectorCountAsync)
            .Where(x => x.Id == id)
            .ExecuteCommandAsync();
    }

    /// <summary>
    /// 创建量化任务
    /// </summary>
    /// <param name="wikiId"></param>
    /// <param name="wikiDetailId"></param>
    /// <param name="remark"></param>
    public async Task<string> CreateQuantizationListAsync(string wikiId, string wikiDetailId, string remark)
    {
        QuantizedList entity = await repository.GetDb().Queryable<QuantizedList>()
            .FirstAsync(x => x.WikiId.Equals(wikiId) && x.WikiDetailId.Equals(wikiDetailId));
        if (entity != null)
        {
            await repository.GetDb().Updateable<QuantizedList>()
                .SetColumns(x => x.State == 1)
                .SetColumns(x => x.Remark == remark)
                .SetColumns(x => x.ModificationTime == DateTime.Now)
                .SetColumns(x => x.ProcessTime == null)
                .Where(x => x.Id == entity.Id)
                .ExecuteCommandAsync();
        }
        else
        {
            entity = await repository.GetDb().Insertable<QuantizedList>(new QuantizedList()
            {
                WikiId = wikiId,
                WikiDetailId = wikiDetailId,
                Remark = remark,
                CreationTime = DateTime.Now,
                State = 1,
            }).ExecuteReturnEntityAsync();
        }

        return entity.Id;
    }

    /// <summary>
    /// 完成量化任务
    /// </summary>
    /// <param name="id"></param>
    /// <param name="remark"></param>
    /// <param name="state">状态 1待处理  2 成功 3 失败</param>
    public async Task CompleteQuantizationListAsync(string id, string remark, int state)
    {
        await repository.GetDb().Updateable<QuantizedList>()
            .SetColumns(x => x.State == state)
            .SetColumns(x => x.Remark == remark)
            .SetColumns(x => x.ModificationTime == DateTime.Now)
            .SetColumns(x => x.ProcessTime == DateTime.Now)
            .Where(x => x.Id == id)
            .ExecuteCommandAsync();
    }


    public async Task<(int total, List<WikiDetailPageListResponse> list)> GetPageListAsync(
        WikiDetailPageRequest request)
    {
        var data = await repository.GetPageListAsync(
            x => ((request.State < 0) || x.State == request.State) && x.WikiId.Equals(request.WikiId),
            request.Page,
            request.PageSize, x => x.CreationTime, OrderByType.Desc);
        List<WikiDetailPageListResponse> result = data.list.Adapt<List<WikiDetailPageListResponse>>();
        if (result.Count > 0)
        {
            var wikiDetailIds = data.list.Select(x => x.Id).ToList();
            var entity = await repository.GetDb().Queryable<QuantizedList>()
                .Where(x => x.WikiId.Equals(request.WikiId) && wikiDetailIds.Contains(x.WikiDetailId)).ToListAsync();
            
            foreach (var wikiDetail in result)
            {
                QuantizedList firstOrDefault = entity.FirstOrDefault(x => x.WikiDetailId.Equals(wikiDetail.Id));
                if (firstOrDefault != null)
                {
                    wikiDetail.Remark = firstOrDefault.Remark;
                }
            }
        }

        return (data.total, result);
    }

    public async Task<List<string>> GetWikiDetailIdsByWikiId(string wikiId)
    {
        return (await repository.GetListAsync(x => x.WikiId.Equals(wikiId) && x.IsDeleted == false))
            .Select(x => x.Id).ToList();
    }

    /// <summary>
    /// 查询向量内容数量
    /// </summary>
    /// <param name="wikiDetailId"></param>
    /// <param name="page"></param>
    /// <param name="size"></param>
    /// <returns></returns>
    public async Task<int> GetWikiDetailVectorCountAsync(string wikiDetailId)
    {
        var detail = await repository.GetByIdAsync(wikiDetailId);
        // 获取 wiki 
        var wiki = await repository.GetDb().Queryable<Domain.Domain.Wikis.Aggregates.Wiki>()
            .InSingleAsync(detail.WikiId);
        if (wiki.VectorType.Equals(VectorType.Mem0))
        {
            return await GetWikiDetailVectorQuantityCountByMem0Async(wiki, detail);
        }

        var setting = await wikiSettingService.GetCommonSetting();
        var memoryServerless = wikiMemoryService.CreateMemoryServerless(wiki.EmbeddingModel, wiki.Model, setting);
        var memoryDbs = memoryServerless.Orchestrator.GetMemoryDbs();
        int total = 0;

        int page = 1;
        int size = 100000;
        foreach (IMemoryDb memoryDb in memoryDbs)
        {
            var limit = page * size;
            if (limit < 10) limit = 10;
            var filter = new MemoryFilter().ByDocument(wikiDetailId);
            await foreach (var item in memoryDb.GetListAsync(OpenAIOption.WikiVectorCollectionName, new List<MemoryFilter>
                           {
                               filter
                           }, limit, true))
            {
                if (item == null || string.IsNullOrEmpty(item.Id))
                {
                    // 证明没有数据了
                    // 直接就退出吧
                    break;
                }

                total++;
            }
        }

        return total;
    }


    /// <summary>
    /// 查询向量内容列表-分页
    /// </summary>
    /// <param name="wikiDetailId"></param>
    /// <param name="page"></param>
    /// <param name="size"></param>
    /// <returns></returns>
    public async Task<(int total, List<WikiDetailVectorQuantityResponse> list)> GetWikiDetailVectorQuantityAsync(
        string wikiDetailId, int page, int size)
    {
        var detail = await repository.GetByIdAsync(wikiDetailId);
        // 获取 wiki 
        var wiki = await repository.GetDb().Queryable<Domain.Domain.Wikis.Aggregates.Wiki>()
            .InSingleAsync(detail.WikiId);

        if (wiki.VectorType.Equals(VectorType.Mem0))
        {
            return await GetWikiDetailVectorQuantityByMem0Async(wiki, detail, page, size: size);
        }


        var setting = await wikiSettingService.GetCommonSetting();
        var memoryServerless = wikiMemoryService.CreateMemoryServerless(wiki.EmbeddingModel, wiki.Model, setting);
        var memoryDbs = memoryServerless.Orchestrator.GetMemoryDbs();

        var list = new List<WikiDetailVectorQuantityResponse>();
        int total = detail.DataCount;

        foreach (IMemoryDb memoryDb in memoryDbs)
        {
            var limit = page * size;
            if (limit < 10) limit = 10;
            var filter = new MemoryFilter().ByDocument(wikiDetailId);
            var dataSize = 0;
            await foreach (var item in memoryDb.GetListAsync(OpenAIOption.WikiVectorCollectionName, new List<MemoryFilter>
                           {
                               filter
                           }, limit, true))
            {
                dataSize++;
                if (dataSize < size * (page - 1)) continue;

                if (dataSize > size * page) break;

                list.Add(new WikiDetailVectorQuantityResponse
                {
                    Content = item.Payload["text"].ToString() ?? string.Empty,
                    FileId = item.Tags.FirstOrDefault(x => x.Key == "fileId").Value?.FirstOrDefault() ?? string.Empty,
                    Id = item.Id,
                    Index = dataSize,
                    WikiDetailId = item.Tags["wikiDetailId"].FirstOrDefault() ?? string.Empty,
                    Document_Id = item.Tags["__document_id"].FirstOrDefault() ?? string.Empty
                });
            }
        }

        return (total, list);
    }

    /// <summary>
    /// 获取向量分页列表-mem0
    /// </summary>
    /// <param name="wiki"></param>
    /// <param name="wikiDetail"></param>
    /// <param name="page"></param>
    /// <param name="size"></param>
    /// <returns></returns>
    private async Task<(int total, List<WikiDetailVectorQuantityResponse> list)> GetWikiDetailVectorQuantityByMem0Async(
        Domain.Domain.Wikis.Aggregates.Wiki wiki,
        WikiDetail wikiDetail, int page, int size)
    {
        int limit = size * (page);
        mem0.NET.Services.MemoryService memoryService =
            serviceProvider.GetRequiredService<mem0.NET.Services.MemoryService>();
        var result = await memoryService.GetMemoryAll(null, wiki.Id.ToString(), wikiDetail.Id, (uint)limit);
        var responses = new List<WikiDetailVectorQuantityResponse>();
        var index = 1;
        foreach (var item in result)
        {
            // 需要跳过前面的数据
            if (index < size * (page - 1))
            {
                index++;
                continue;
            }

            if (index > size * page) break;

            responses.Add(new WikiDetailVectorQuantityResponse
            {
                Content = item.MetaData["metaData"],
                FileId = item.MetaData.FirstOrDefault(x => x.Key == "fileId").Value ?? string.Empty,
                Id = item.Id.ToString(),
                Index = index++,
                WikiDetailId = item.MetaData["wikiDetailId"] ?? string.Empty,
                Document_Id = item.Id.ToString()
            });
        }

        return (result.Count, responses);
    }

    /// <summary>
    /// 查询向量内容数量
    /// </summary>
    /// <param name="wiki"></param>
    /// <param name="wikiDetail"></param>
    /// <param name="page"></param>
    /// <param name="size"></param>
    /// <returns></returns>
    private async Task<int> GetWikiDetailVectorQuantityCountByMem0Async(
        Domain.Domain.Wikis.Aggregates.Wiki wiki,
        WikiDetail wikiDetail)
    {
        int page = 1;
        int size = 100000;
        int limit = page * size;
        mem0.NET.Services.MemoryService memoryService =
            serviceProvider.GetRequiredService<mem0.NET.Services.MemoryService>();
        var result = await memoryService.GetMemoryAll(null, wiki.Id.ToString(), wikiDetail.Id, (uint)limit);

        var index = 0;
        foreach (var item in result)
        {
            if (item == null || string.IsNullOrEmpty(item.Id.ToString()))
            {
                // 证明没有数据了
                // 直接就退出吧
                break;
            }

            index++;
        }

        return index;
    }

    /// <summary>
    /// 检索匹配向量
    /// </summary>
    /// <param name="wikiId"></param>
    /// <param name="minRelevance"></param>
    /// <param name="search"></param>
    public async Task<SearchVectorQuantityResult> SearchVectorQuantityAsync(string wikiId, double minRelevance,
        string search)
    {
        var stopwatch = Stopwatch.StartNew();
        var wiki = await repository.GetDb().Queryable<Domain.Domain.Wikis.Aggregates.Wiki>().InSingleAsync(wikiId);
        var setting = await wikiSettingService.GetCommonSetting();
        var memoryServerless = wikiMemoryService.CreateMemoryServerless(wiki.EmbeddingModel, wiki.Model, setting);
        // search result 
        var searchResult = await memoryServerless.SearchAsync(search, OpenAIOption.WikiVectorCollectionName,
            new MemoryFilter().ByTag("wikiId", wikiId), minRelevance: minRelevance, limit: 5);

        stopwatch.Stop();
        var searchVectorQuantityResult = new SearchVectorQuantityResult();
        searchVectorQuantityResult.ElapsedTime = stopwatch.ElapsedMilliseconds;

        searchVectorQuantityResult.Result = new List<SearchVectorQuantityResponse>();
        foreach (var resultResult in searchResult.Results)
        {
            searchVectorQuantityResult.Result.AddRange(resultResult.Partitions.Select(partition =>
                {
                    partition.Tags.TryGetValue("fileId", out var data);
                    return new SearchVectorQuantityResponse()
                    {
                        Content = partition.Text,
                        DocumentId = resultResult.DocumentId,
                        Relevance = partition.Relevance,
                        FileId = data?.FirstOrDefault() ?? string.Empty
                    };
                }
            ));
        }

        searchVectorQuantityResult.Result =
            searchVectorQuantityResult.Result.OrderByDescending(x => x.Relevance).ToList();
        return searchVectorQuantityResult;
    }

    /// <summary>
    /// 删除向量
    /// </summary>
    /// <param name="documentId"></param>
    /// <param name="wikiDetailId"></param>
    public async Task RemoveWikiDetailVectorQuantityAsync(string documentId, string wikiDetailId)
    {
        var wiki = await GetWikiByDetailId(wikiDetailId);
        var setting = await wikiSettingService.GetCommonSetting();
        var memoryServerless = wikiMemoryService.CreateMemoryServerless(wiki.EmbeddingModel, wiki.Model, setting);
        await memoryServerless.DeleteDocumentAsync(documentId, OpenAIOption.WikiVectorCollectionName);
    }


    /// <summary>
    /// 删除知识点
    /// </summary>
    /// <param name="id"></param>
    public async Task RemoveWikiDetailAsync(string id)
    {
        var wiki = await GetWikiByDetailId(id);
        if (wiki == null)
        {
            throw new ArgumentException("知识库不存在");
        }

        await repository.DeleteAsync(id);
        var setting = await wikiSettingService.GetCommonSetting();

        try
        {
            if (wiki?.VectorType == VectorType.Mem0)
            {
                mem0.NET.Services.MemoryService memoryService =
                    serviceProvider.GetRequiredService<mem0.NET.Services.MemoryService>();
                await memoryService.DeleteAll(null, wiki.Id, id);
            }
            else
            {
                var memoryServerless =
                    wikiMemoryService.CreateMemoryServerless(wiki.EmbeddingModel, wiki.Model, setting);
                await memoryServerless.DeleteDocumentAsync(id, OpenAIOption.WikiVectorCollectionName);
            }
        }
        catch (Exception e)
        {
            Console.WriteLine(e);
        }
    }

    private async Task<Domain.Domain.Wikis.Aggregates.Wiki> GetWikiByDetailId(string id)
    {
        WikiDetail byIdAsync = await repository.GetByIdAsync(id);
        string wikiId = byIdAsync.WikiId;
        var wiki = await repository.GetDb().Queryable<Domain.Domain.Wikis.Aggregates.Wiki>().InSingleAsync(wikiId);
        return wiki;
    }

    public async Task<WikiDetail> GetByIdAsync(string id)
    {
        return await repository.GetByIdAsync(id);
    }
}